{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8f5c6f19",
   "metadata": {},
   "source": [
    "# `LayerNorm` 前端支持"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87b9eaf9",
   "metadata": {},
   "source": [
    "## {func}`~tvm.relax.frontend.torch.from_exported_program`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2953dabc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.export import export\n",
    "import tvm\n",
    "from tvm import relax\n",
    "import tvm.testing\n",
    "from tvm.script import ir as I\n",
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T\n",
    "from tvm.relax.frontend.torch import from_exported_program\n",
    "class M(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.ln = torch.nn.RMSNorm((32))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.ln(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fcf97fe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[torch.onnx] Obtain model graph for `M([...]` with `torch.export.export(..., strict=False)`...\n",
      "[torch.onnx] Obtain model graph for `M([...]` with `torch.export.export(..., strict=False)`... ✅\n",
      "[torch.onnx] Dumping ExportedProgram because `dump_exported_program=True`...\n",
      "[torch.onnx] ExportedProgram has been saved to 'onnx_export_2025-09-12_17-11-16-671771.pt2'.\n",
      "[torch.onnx] Run decomposition...\n",
      "[torch.onnx] Run decomposition... ✅\n",
      "[torch.onnx] Translate the graph into ONNX...\n",
      "[torch.onnx] Translate the graph into ONNX... ✅\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ONNXProgram(\n",
       "    model=\n",
       "        <\n",
       "            ir_version=10,\n",
       "            opset_imports={'': 23},\n",
       "            producer_name='pytorch',\n",
       "            producer_version='2.8.0+cu128',\n",
       "            domain=None,\n",
       "            model_version=None,\n",
       "        >\n",
       "        graph(\n",
       "            name=main_graph,\n",
       "            inputs=(\n",
       "                %\"x\"<FLOAT,[1,3,32,32]>\n",
       "            ),\n",
       "            outputs=(\n",
       "                %\"type_as\"<FLOAT,[1,3,32,32]>\n",
       "            ),\n",
       "            initializers=(\n",
       "                %\"ln.weight\"<FLOAT,[32]>{TorchTensor(...)},\n",
       "                %\"val_0\"<FLOAT,[]>{Tensor<FLOAT,[]>(array(2., dtype=float32), name='val_0')},\n",
       "                %\"val_3\"<INT64,[1]>{Tensor<INT64,[1]>(array([3]), name='val_3')},\n",
       "                %\"val_4\"<FLOAT,[]>{Tensor<FLOAT,[]>(array(1.1920929e-07, dtype=float32), name='val_4')}\n",
       "            ),\n",
       "        ) {\n",
       "            0 |  # node_pow_1\n",
       "                 %\"pow_1\"<FLOAT,[1,3,32,32]> ⬅️ ::Pow(%\"x\", %\"val_0\"{2.0})\n",
       "            1 |  # node_mean\n",
       "                 %\"mean\"<FLOAT,[1,3,32,1]> ⬅️ ::ReduceMean(%\"pow_1\", %\"val_3\"{[3]}) {keepdims=True, noop_with_empty_axes=0}\n",
       "            2 |  # node_add\n",
       "                 %\"add\"<FLOAT,[1,3,32,1]> ⬅️ ::Add(%\"mean\", %\"val_4\"{1.1920928955078125e-07})\n",
       "            3 |  # node_Sqrt_5\n",
       "                 %\"val_5\"<FLOAT,[1,3,32,1]> ⬅️ ::Sqrt(%\"add\")\n",
       "            4 |  # node_rsqrt\n",
       "                 %\"rsqrt\"<FLOAT,[1,3,32,1]> ⬅️ ::Reciprocal(%\"val_5\")\n",
       "            5 |  # node_mul\n",
       "                 %\"mul\"<FLOAT,[1,3,32,32]> ⬅️ ::Mul(%\"x\", %\"rsqrt\")\n",
       "            6 |  # node_type_as\n",
       "                 %\"type_as\"<FLOAT,[1,3,32,32]> ⬅️ ::Mul(%\"mul\", %\"ln.weight\"{...})\n",
       "            return %\"type_as\"<FLOAT,[1,3,32,32]>\n",
       "        }\n",
       "\n",
       "\n",
       "    ,\n",
       "    exported_program=\n",
       "        ExportedProgram:\n",
       "            class GraphModule(torch.nn.Module):\n",
       "                def forward(self, p_ln_weight: \"f32[32]\", x: \"f32[1, 3, 32, 32]\"):\n",
       "                     # File: /media/pc/data/lxw/envs/anaconda3a/envs/py313/lib/python3.13/site-packages/torch/nn/modules/normalization.py:402 in forward, code: return F.rms_norm(x, self.normalized_shape, self.weight, self.eps)\n",
       "                    pow_1: \"f32[1, 3, 32, 32]\" = torch.ops.aten.pow.Tensor_Scalar(x, 2)\n",
       "                    mean: \"f32[1, 3, 32, 1]\" = torch.ops.aten.mean.dim(pow_1, [3], True);  pow_1 = None\n",
       "                    add: \"f32[1, 3, 32, 1]\" = torch.ops.aten.add.Scalar(mean, 1.1920928955078125e-07);  mean = None\n",
       "                    rsqrt: \"f32[1, 3, 32, 1]\" = torch.ops.aten.rsqrt.default(add);  add = None\n",
       "                    mul: \"f32[1, 3, 32, 32]\" = torch.ops.aten.mul.Tensor(x, rsqrt);  rsqrt = None\n",
       "                    mul_1: \"f32[1, 3, 32, 32]\" = torch.ops.aten.mul.Tensor(mul, p_ln_weight);  mul = p_ln_weight = None\n",
       "                    type_as: \"f32[1, 3, 32, 32]\" = torch.ops.aten.type_as.default(mul_1, x);  mul_1 = x = None\n",
       "                    return (type_as,)\n",
       "            \n",
       "        Graph signature: \n",
       "            # inputs\n",
       "            p_ln_weight: PARAMETER target='ln.weight'\n",
       "            x: USER_INPUT\n",
       "    \n",
       "            # outputs\n",
       "            type_as: USER_OUTPUT\n",
       "    \n",
       "        Range constraints: {}\n",
       "\n",
       ")"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_shape = [1, 3, 32, 32]\n",
    "model = M().eval()\n",
    "input_data = torch.randn(input_shape)\n",
    "torch.onnx.export(\n",
    "    model, \n",
    "    input_data,\n",
    "    \".temp/test.onnx\", \n",
    "    input_names=[\"x\"],\n",
    "    dynamo=True,\n",
    "    opset_version=23,\n",
    "    do_constant_folding=True,\n",
    "    export_params=True,\n",
    "    dump_exported_program=True,\n",
    "    keep_initializers_as_inputs=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8652beaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)\n",
    "\n",
    "model = LayerNorm()\n",
    "exported_program = export(model, args=example_args,)\n",
    "mod = from_exported_program(exported_program)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55a1debe",
   "metadata": {},
   "source": [
    "## {class}`~tvm.relax.frontend.nn.modules.LayerNorm`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaf53516",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.relax.frontend.nn import core, modules, spec\n",
    "mod = modules.LayerNorm(8)\n",
    "tvm_mod, params = mod.export_tvm(\n",
    "    spec={\"forward\": {\"x\": spec.Tensor((2, 4, 8), \"float32\")}}, debug=True\n",
    ")\n",
    "tvm_mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0526679b",
   "metadata": {},
   "source": [
    "## {func}`~tvm.relax.frontend.torch.from_fx`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c4b8108",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm import relax\n",
    "import tvm.testing\n",
    "from tvm.script import ir as I\n",
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T\n",
    "from tvm.relax.frontend import detach_params\n",
    "from tvm.relax.frontend.torch import from_fx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae81b996",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNorm(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.ln = torch.nn.LayerNorm((10, 10))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.ln(x)\n",
    "\n",
    "input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "\n",
    "torch_model = LayerNorm()\n",
    "graph_model = torch.fx.symbolic_trace(torch_model)\n",
    "with torch.no_grad():\n",
    "    mod = from_fx(graph_model, input_info)\n",
    "    mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b8c3d16",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNorm(torch.nn.Module):\n",
    "    def __init__(self, shape):\n",
    "        super().__init__()\n",
    "        self.weight = torch.nn.Parameter(torch.ones(shape))\n",
    "        self.bias = torch.nn.Parameter(torch.zeros(shape))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return torch.nn.functional.layer_norm(\n",
    "            x, self.weight.shape, self.weight, self.bias, 1e-5\n",
    "        )\n",
    "\n",
    "torch_model = LayerNorm((10, 10))\n",
    "graph_model = torch.fx.symbolic_trace(torch_model)\n",
    "with torch.no_grad():\n",
    "    mod = from_fx(graph_model, input_info)\n",
    "    mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d038bd2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNorm2(torch.nn.Module):\n",
    "    def __init__(self, shape):\n",
    "        super().__init__()\n",
    "        self.shape = shape\n",
    "        self.weight = None\n",
    "        self.bias = None\n",
    "\n",
    "    def forward(self, input):\n",
    "        return torch.nn.functional.layer_norm(input, self.shape, self.weight, self.bias, 1e-5)\n",
    "\n",
    "torch_model = LayerNorm2((10, 10))\n",
    "graph_model = torch.fx.symbolic_trace(torch_model)\n",
    "with torch.no_grad():\n",
    "    mod = from_fx(graph_model, input_info)\n",
    "    mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "976243fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNorm3(torch.nn.Module):\n",
    "    def __init__(self, shape):\n",
    "        super().__init__()\n",
    "        self.shape = shape\n",
    "        self.weight = torch.nn.Parameter(torch.ones(shape))\n",
    "        self.bias = torch.nn.Parameter(torch.zeros(shape))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return torch.nn.functional.layer_norm(x, self.shape, self.weight, self.bias, 1e-5)\n",
    "\n",
    "torch_model = LayerNorm3((10, 10))\n",
    "graph_model = torch.fx.symbolic_trace(torch_model)\n",
    "with torch.no_grad():\n",
    "    mod = from_fx(graph_model, input_info)\n",
    "    mod.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
